from typing import List, Any, Optional, TypeVar, Union
import numpy as np
import time
import logging

T = TypeVar("T")


class NamespaceDict(dict):
    """
    A dictionary that supports both attribute-style and dictionary-style access.

    This class extends the standard dictionary by allowing keys to be accessed as attributes.
    It automatically converts nested dictionaries and lists containing dictionaries to
    NamespaceDict instances recursively, providing a convenient way to access deeply
    nested configurations.

    Key Features:
        - Access dictionary items as attributes: `config.section.option` instead of `config['section']['option']`
        - Set values using either attribute or dictionary syntax
        - Delete keys/attributes using either syntax
        - Automatic conversion of nested dictionaries to NamespaceDict instances
        - Beautiful tree-like string representation for complex nested structures

    Special Considerations:
        - WARNING: Using dictionary method names as keys (e.g., 'items', 'keys', 'update') may cause
          unexpected behavior since attribute access would return the value rather than the method.
        - To safely access dictionary methods when such keys exist, use the safe_* methods provided
          (e.g., safe_items(), safe_keys()) or use dict.method(instance) directly.

    Typical Use Cases:
        - Configuration objects
        - Data structures where attribute-style access is more readable
        - JSON/YAML loaded data that needs attribute access

    Examples:
        Basic usage:
            >>> config = NamespaceDict(server={"host": "localhost", "port": 8080})
            >>> config.server.host
            'localhost'
            >>> config.server.port = 9000
            >>> config.server.port
            9000
            >>> config["server"]["port"]
            9000

        Nested structures:
            >>> settings = NamespaceDict()
            >>> settings.database = {"url": "postgresql://user:pass@localhost/db", "timeout": 30}
            >>> settings.database.timeout = 60
            >>> settings.database.timeout
            60

        Deletion:
            >>> d = NamespaceDict(a=1, b=2)
            >>> del d.a
            >>> hasattr(d, 'a')
            False
            >>> 'a' in d
            False

        Key name conflicts:
            >>> d = NamespaceDict(items=1, keys=2, update=3)
            >>> d.items  # Returns 1, not the items() method
            1
            >>> list(d.safe_items())  # Use safe_items() to access the method
            [('items', 1), ('keys', 2), ('update', 3)]
    """

    # Class-level list of reserved method names that might conflict with attribute access
    _RESERVED_NAMES = {
        "items",
        "keys",
        "values",
        "update",
        "pop",
        "popitem",
        "clear",
        "copy",
        "get",
        "setdefault",
        "fromkeys",
    }

    def __init__(self, *args, **kwargs) -> None:
        """
        Initialize the NamespaceDict with the given items.

        Converts all dictionary values to NamespaceDict instances recursively.

        Args:
            *args: Variable length argument list (same as dict)
            **kwargs: Arbitrary keyword arguments (same as dict)

        Examples:
            >>> d = NamespaceDict({"a": 1, "b": {"c": 2}})
            >>> isinstance(d.b, NamespaceDict)
            True
            >>> d = NamespaceDict(x=1, y=[{"z": 2}])
            >>> isinstance(d.y[0], NamespaceDict)
            True

        Note:
            If you use reserved method names as keys (e.g., 'items', 'keys'), they will
            still work as attributes, but you'll need to use safe_* methods to access the
            original dictionary functionality.
        """
        super().__init__(*args, **kwargs)
        # Convert initial values and update both dict and __dict__
        for key, value in list(dict.items(self)):
            converted = self._convert_value(value)
            super().__setitem__(key, converted)  # Update dict part
            self.__dict__[key] = converted  # Update __dict__

        # Warn if any reserved names are used as keys
        conflicting_keys = [
            key for key in dict.keys(self) if key in self._RESERVED_NAMES
        ]
        if conflicting_keys:
            import warnings

            warnings.warn(
                f"Keys {conflicting_keys} conflict with built-in dictionary methods. "
                f"Use safe_* methods (e.g., safe_{conflicting_keys[0]}()) "
                f"to access the dictionary methods.",
                UserWarning,
            )

    def _convert_value(self, value: Any) -> Any:
        """
        Convert dictionaries to NamespaceDict and handle nested structures.

        This is a recursive method that processes nested dictionaries and collections
        that might contain dictionaries.

        Args:
            value: Value to potentially convert

        Returns:
            Converted value:
                - Dictionaries become NamespaceDict objects
                - Lists/tuples have their dictionary elements converted
                - Other values remain unchanged

        Examples:
            >>> d = NamespaceDict()
            >>> converted = d._convert_value({"a": 1, "b": [{"c": 2}]})
            >>> isinstance(converted, NamespaceDict)
            True
            >>> isinstance(converted.b[0], NamespaceDict)
            True
        """
        if isinstance(value, dict) and not isinstance(value, NamespaceDict):
            return NamespaceDict(value)
        elif isinstance(value, (list, tuple)):
            return type(value)(self._convert_value(x) for x in value)
        return value

    def __getattr__(self, key: str) -> Any:
        """
        Access items as attributes.

        This method is called when an attribute is accessed that doesn't exist
        in the instance's __dict__. It allows attribute-style access to dictionary keys.

        Args:
            key: Attribute name to access

        Returns:
            The value associated with the key

        Raises:
            AttributeError: If the key doesn't exist in the dictionary

        Examples:
            >>> d = NamespaceDict(x=1)
            >>> d.x
            1
            >>> d.y  # Raises AttributeError
            AttributeError: 'NamespaceDict' object has no attribute 'y'
        """
        try:
            return self.__dict__[key]
        except KeyError:
            raise AttributeError(f"'NamespaceDict' object has no attribute '{key}'")

    def __setattr__(self, name: str, value: Any) -> None:
        """
        Set attribute value, handling nested structures.

        Sets both the dictionary item and the actual attribute,
        converting nested dictionaries to NamespaceDict objects.

        Args:
            name: Attribute name
            value: Value to set

        Examples:
            >>> d = NamespaceDict()
            >>> d.x = {"y": 1}
            >>> isinstance(d.x, NamespaceDict)
            True
            >>> d.x.y
            1
        """
        value = self._convert_value(value)
        self.__dict__[name] = value
        self[name] = value

    def __getitem__(self, key: str) -> Any:
        """
        Access items as dictionary keys.
        Provides standard dictionary access syntax.

        Args:
            key: Dictionary key to access

        Returns:
            The value associated with the key

        Raises:
            KeyError: If the key doesn't exist in the dictionary

        Examples:
            >>> d = NamespaceDict(x=1)
            >>> d["x"]
            1
            >>> d["y"]  # Raises KeyError
            KeyError: "'NamespaceDict' object has no key 'y'"
        """
        try:
            return self.__dict__[key]
        except KeyError:
            raise KeyError(f"'NamespaceDict' object has no key '{key}'")

    def __setitem__(self, key: str, value: Any) -> None:
        """
        Set dictionary item, handling nested structures.

        Sets both the dictionary item and the corresponding attribute,
        converting nested dictionaries to NamespaceDict objects.

        Args:
            key: Dictionary key
            value: Value to set

        Examples:
            >>> d = NamespaceDict()
            >>> d["config"] = {"debug": True}
            >>> d.config.debug
            True
        """
        value = self._convert_value(value)
        super().__setitem__(key, value)
        self.__dict__[key] = value

    def __delattr__(self, key: str) -> None:
        """
        Delete an attribute and corresponding dictionary item.

        Removes both the attribute and the dictionary item with the given key.

        Args:
            key: Attribute name to delete

        Raises:
            AttributeError: If the key doesn't exist

        Examples:
            >>> d = NamespaceDict(x=1)
            >>> del d.x
            >>> hasattr(d, 'x')
            False
        """
        try:
            super().__delitem__(key)
            del self.__dict__[key]
        except KeyError:
            raise AttributeError(f"'NamespaceDict' object has no attribute '{key}'")

    def __delitem__(self, key: str) -> None:
        """
        Delete dictionary item and corresponding attribute.

        Removes both the dictionary item and the attribute with the given key.

        Args:
            key: Dictionary key to delete

        Raises:
            KeyError: If the key doesn't exist in the dictionary

        Examples:
            >>> d = NamespaceDict(x=1)
            >>> del d["x"]
            >>> "x" in d
            False
        """
        super().__delitem__(key)
        try:
            del self.__dict__[key]
        except KeyError:
            pass  # Already deleted or never was an attribute

    def update(self, *args, **kwargs) -> None:
        """
        Update the dictionary and attributes with new key-value pairs.

        Updates both dictionary items and attributes, converting nested dictionaries to
        NamespaceDict objects automatically.

        Args:
            *args: Variable length argument list (same as dict.update)
            **kwargs: Arbitrary keyword arguments (same as dict.update)

        Examples:
            >>> d = NamespaceDict(a=1)
            >>> d.update({"b": 2, "c": {"d": 3}})
            >>> d.a, d.b, d.c.d
            (1, 2, 3)
            >>> d.update(e=4, f={"g": 5})
            >>> d.e, d.f.g
            (4, 5)
        """
        items = dict(*args, **kwargs)
        for key, value in items.items():
            self[key] = value  # Uses __setitem__ which handles conversion

    def pop(self, key: str, default: Optional[T] = None) -> Union[Any, T]:
        """
        Pop a key from the dictionary and attributes.

        Removes the key from both the dictionary and attributes, returning the associated value.

        Args:
            key: Key to pop
            default: Default value if key doesn't exist

        Returns:
            The value associated with the key, or default if provided and key doesn't exist

        Raises:
            KeyError: If the key doesn't exist and no default is provided

        Examples:
            >>> d = NamespaceDict(a=1, b=2)
            >>> d.pop("a")
            1
            >>> "a" in d
            False
            >>> d.pop("c", 3)  # Key doesn't exist, return default
            3
        """
        try:
            val = super().pop(key)
            try:
                del self.__dict__[key]
            except KeyError:
                pass
            return val
        except KeyError:
            if default is not None:
                return default
            raise KeyError(key)

    def __contains__(self, key: object) -> bool:
        """
        Check if a key exists in the dictionary.

        Args:
            key: Key to check

        Returns:
            True if the key exists, False otherwise

        Examples:
            >>> d = NamespaceDict(a=1)
            >>> "a" in d
            True
            >>> "b" in d
            False
        """
        return key in self.__dict__

    def __eq__(self, other: object) -> bool:
        """
        Check equality with another object.

        Two NamespaceDict instances are equal if their underlying dictionaries are equal.
        A NamespaceDict is equal to a regular dict if the underlying dictionary is equal to it.

        Args:
            other: Object to compare with

        Returns:
            True if equal, False otherwise

        Examples:
            >>> d1 = NamespaceDict(a=1, b=2)
            >>> d2 = NamespaceDict(a=1, b=2)
            >>> d1 == d2
            True
            >>> d1 == {"a": 1, "b": 2}
            True
        """
        if not isinstance(other, (NamespaceDict, dict)):
            return False
        if isinstance(other, NamespaceDict):
            return dict(self) == dict(other)
        return dict(self) == other

    def __repr__(self) -> str:
        """
        Developer-oriented string representation of the NamespaceDict.

        Provides a representation that could be used to recreate the object,
        with improved formatting for better readability.

        Returns:
            A string representation of the object that could be used to recreate it

        Examples:
            >>> d = NamespaceDict(a=1, b=2)
            >>> repr(d)
            "NamespaceDict(a=1, b=2)"

            >>> d = NamespaceDict(a=1, b=2, c=3, d=4)  # More items will use multiline format
            >>> print(repr(d))
            NamespaceDict(
                a=1,
                b=2,
                c=3,
                d=4
            )
        """
        items = []
        for k, v in self.items():
            items.append(f"{k}={repr(v)}")

        if len(items) <= 3:  # 如果项目少，保持在一行显示
            return f"{type(self).__name__}({', '.join(items)})"
        else:  # 否则使用多行格式
            joined_items = ",\n    ".join(items)
            return f"{type(self).__name__}(\n    {joined_items}\n)"

    def __str__(self) -> str:
        """
        User-friendly formatted string representation of the NamespaceDict.

        Returns a beautifully formatted tree-like structure showing the hierarchy
        of nested objects, making it easy to visualize complex configurations.

        Returns:
            A nicely formatted string representation

        Examples:
            >>> config = NamespaceDict(
            ...     database={"host": "localhost", "port": 5432},
            ...     server={"debug": True, "workers": 4}
            ... )
            >>> print(config)
            database =
            ├── host = 'localhost'
            └── port = 5432
            server =
            ├── debug = True
            └── workers = 4
        """
        return self._format_dict()

    def _format_dict(self, indent=0, is_last=True, prefix_str=""):
        """
        Format dictionary with beautiful tree-like structure.

        Creates a hierarchical tree visualization of the dictionary structure,
        with appropriate indentation and connection lines.

        Args:
            indent: Current indentation level
            is_last: Whether this is the last item in parent
            prefix_str: Prefix string for current line

        Returns:
            Formatted string representation of the dictionary structure

        Note:
            This is a recursive method used internally by __str__
        """
        if not self:
            return ""

        result = []
        # 使用 dict.items 而不是 self.items() 以避免与属性冲突
        items = list(dict.items(self))

        # Calculate maximum key length for alignment
        max_key_len = max((len(str(k)) for k in dict.keys(self)), default=0)

        for i, (k, v) in enumerate(items):
            is_last_item = i == len(items) - 1

            # Determine current line prefix
            if indent == 0:  # Top-level elements
                curr_prefix = ""
                next_prefix = "    "
            else:
                curr_prefix = prefix_str + ("└── " if is_last else "├── ")
                next_prefix = prefix_str + ("    " if is_last else "│   ")

            # Format key name with alignment
            key_str = f"{k:<{max_key_len}}"

            if isinstance(v, NamespaceDict):
                # Handle nested dictionaries
                result.append(f"{curr_prefix}{key_str} = ")
                nested_str = v._format_dict(indent + 1, is_last_item, next_prefix)
                if nested_str:
                    result.append(nested_str)

            elif isinstance(v, (list, tuple)):
                # Handle lists or tuples
                if any(isinstance(x, NamespaceDict) for x in v):
                    result.append(f"{curr_prefix}{key_str} = [")
                    for j, item in enumerate(v):
                        is_last_list_item = j == len(v) - 1
                        list_prefix = next_prefix + (
                            "└── " if is_last_list_item else "├── "
                        )

                        if isinstance(item, NamespaceDict):
                            result.append(
                                f"{next_prefix}{'└── ' if is_last_list_item else '├── '}Item {j}:"
                            )
                            nested_str = item._format_dict(
                                indent + 2,
                                is_last_list_item,
                                next_prefix + ("    " if is_last_list_item else "│   "),
                            )
                            if nested_str:
                                result.append(nested_str)
                        else:
                            result.append(f"{list_prefix}{repr(item)}")

                    result.append(f"{next_prefix}]")
                else:
                    # Simple list representation
                    result.append(f"{curr_prefix}{key_str} = {repr(v)}")
            else:
                # Simple value representation
                result.append(f"{curr_prefix}{key_str} = {repr(v)}")

        return "\n".join(result)

    # Add safe access methods for all reserved methods
    def safe_items(self):
        """Safely access the items() method even if 'items' is used as a key."""
        return dict.items(self)

    def safe_keys(self):
        """Safely access the keys() method even if 'keys' is used as a key."""
        return dict.keys(self)

    def safe_values(self):
        """Safely access the values() method even if 'values' is used as a key."""
        return dict.values(self)

    def safe_get(self, key, default=None):
        """Safely access the get() method even if 'get' is used as a key."""
        return dict.get(self, key, default)

    def safe_pop(self, key, default=None):
        """Safely access the pop() method even if 'pop' is used as a key.

        Note: This affects both dictionary entries and attributes.
        """
        try:
            val = dict.pop(self, key)
            try:
                del self.__dict__[key]
            except KeyError:
                pass
            return val
        except KeyError:
            if default is not None:
                return default
            raise KeyError(key)

    def safe_update(self, *args, **kwargs):
        """Safely access the update() method even if 'update' is used as a key."""
        items = dict(*args, **kwargs)
        for key, value in items.items():
            self[key] = value

    def get_dict_item(self, key, default=None):
        """
        Get a dictionary item safely, regardless of attribute name conflicts.

        Args:
            key: The key to retrieve
            default: Value to return if key doesn't exist

        Returns:
            The value associated with the key or default if not found
        """
        return dict.get(self, key, default)

    def get_dict_items(self):
        """
        Get all items in the dictionary safely, regardless of attribute name conflicts.

        Returns:
            A view of the dictionary's items
        """
        return dict.items(self)


def time_execution(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()  # Start timer
        result = func(*args, **kwargs)  # Call the original function
        end_time = time.time()  # End timer
        elapsed_time = end_time - start_time
        print(f"Time taken for {func.__name__:25s}: {elapsed_time:.2f} seconds")
        return result

    return wrapper


def count_parameters(model, trainable=False):
    if trainable:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())


def accuracy_dil(y_pred, y_true, nb_old, inc_cls, class_num):
    assert len(y_pred) == len(y_true), "Data length error."
    all_acc = {}
    all_acc["total"] = np.around(
        (y_pred % class_num == y_true % class_num).sum() * 100 / len(y_true),
        decimals=2,
    )

    # Grouped accuracy
    for class_id in range(0, np.max(y_true), inc_cls):
        idxes = np.where(
            np.logical_and(y_true >= class_id, y_true < class_id + inc_cls)
        )[0]
        label = "{}-{}".format(
            str(class_id).rjust(2, "0"),
            str(class_id + inc_cls - 1).rjust(2, "0"),
        )
        all_acc[label] = np.around(
            ((y_pred[idxes] % class_num) == (y_true[idxes] % class_num)).sum()
            * 100
            / len(idxes),
            decimals=2,
        )

    # Old accuracy
    idxes = np.where(y_true < nb_old)[0]
    all_acc["old"] = (
        0
        if len(idxes) == 0
        else np.around(
            ((y_pred[idxes] % class_num) == (y_true[idxes] % class_num)).sum()
            * 100
            / len(idxes),
            decimals=2,
        )
    )

    # New accuracy
    idxes = np.where(y_true >= nb_old)[0]
    all_acc["new"] = np.around(
        ((y_pred[idxes] % class_num) == (y_true[idxes] % class_num)).sum()
        * 100
        / len(idxes),
        decimals=2,
    )

    return all_acc


def accuracy_cil(y_pred, y_true, nb_old, increment):
    assert len(y_pred) == len(y_true), "Data length error."
    all_acc = {}
    all_acc["total"] = np.around(
        (y_pred == y_true).sum() * 100 / len(y_true), decimals=2
    )

    # FIXME: This is a temporary solution to load the ViT models for fine-tuning
    increment = increment if increment > 0 else np.max(y_true)
    # Grouped accuracy
    for class_id in range(0, np.max(y_true), increment):
        idxes = np.where(
            np.logical_and(y_true >= class_id, y_true < class_id + increment)
        )[0]
        label = "{}-{}".format(
            str(class_id).rjust(2, "0"),
            str(class_id + increment - 1).rjust(2, "0"),
        )
        all_acc[label] = np.around(
            (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
            decimals=2,
        )

    # Old accuracy
    idxes = np.where(y_true < nb_old)[0]

    all_acc["old"] = (
        0
        if len(idxes) == 0
        else np.around(
            (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
            decimals=2,
        )
    )

    # New accuracy
    idxes = np.where(y_true >= nb_old)[0]
    all_acc["new"] = np.around(
        (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
        decimals=2,
    )

    return all_acc


def split_images_labels(imgs):
    # split trainset.imgs in ImageFolder
    images = []
    labels = []
    for item in imgs:
        images.append(item[0])
        labels.append(item[1])

    return np.array(images), np.array(labels)


class AccMeter:
    """Computes and stores average and current values for an arbitrary number of metrics (e.g., losses)."""

    def __init__(self, names, fmt=".4f", tail=""):
        """Initialize with names of metrics."""
        self.names = (
            names  # List of metric names (e.g., ['loss1', 'loss2', 'accuracy'])
        )
        self.fmt = fmt
        self.tail = tail
        self.reset()

    def add(self, names: List):
        """Add new metrics to the meter."""
        for name in names:
            if name in self.names:
                continue
            self.names.append(name)
            self.avg[name] = 0.0
            self.sum[name] = 0.0
            self.count[name] = 0.0

    def reset(self):
        """Reset all meters."""
        self.avg = {name: 0.0 for name in self.names}
        self.sum = {name: 0.0 for name in self.names}
        self.count = {name: 0.0 for name in self.names}

    def update(self, dict_vals):
        """Update all metrics with new values."""
        for name, val in dict_vals.items():
            self.sum[name] += val["correct"]
            self.count[name] += val["total"]
            self.avg[name] = (self.sum[name] / self.count[name]) * 100

    def __str__(self):
        """Return formatted string for displaying all metrics."""
        metrics_str = []
        for name in self.names:
            fmtstr = f"{name:10s}: {{avg:{self.fmt}}}" + self.tail
            metrics_str.append(fmtstr.format(avg=self.avg[name], name=name))
        return "  ".join(metrics_str)


class LossMeter(AccMeter):
    def update(self, dict_vals):
        for name, val in dict_vals.items():
            self.sum[name] += val
            self.count[name] += 1
            self.avg[name] = self.sum[name] / self.count[name]


class MeterProcessor:
    """Handles the display of metric progress across epochs."""

    def __init__(self, name, epochs, meters, prefix=""):
        self.name, self.epoch_fmt, self.meters, self.prefix = (
            name,
            self.__epoch_fmtstr(epochs),
            meters,
            prefix,
        )

    def display(self, epoch, *args, **kwargs):
        """Display metrics for the current epoch."""
        entries = [f"{self.prefix}{self.epoch_fmt.format(epoch)}"] + [
            str(meter) for meter in self.meters
        ]
        logging.info("  ".join(entries))

    def __epoch_fmtstr(self, epochs):
        """Generate format string for displaying epoch information."""
        num_digits = len(str(int(epochs)))
        return f"[{{:0{num_digits}d}}/{epochs}]"


class TrainMP(MeterProcessor):
    """Extended MeterProcessor for displaying training-specific metrics (including multiple losses)."""

    def __init__(self, name, epochs, meters, prefix=""):
        super().__init__(name, epochs, meters, prefix)
        self.timer = None

    def reset_timer(self):
        """Start the timer."""
        self.timer = time.time()

    def log_timer(self):
        """Stop the timer and return elapsed time."""
        if self.timer is not None:
            elapsed_time = time.time() - self.timer
            return elapsed_time
        return None

    def display(self, epoch, lr: List):
        """Display training metrics along with timer and learning rate."""
        elapsed_time = self.log_timer()
        timer_str = f"[{elapsed_time:.2f} seconds]" if elapsed_time else "[No Timer]"
        lr_str = f"{[f'{i:.3e}' for i in lr]}\n"
        entries = [
            f"{self.prefix}{self.epoch_fmt.format(epoch)} {timer_str} {lr_str}"
        ] + [str(meter) for meter in self.meters]
        logging.info("  ".join(entries))
